#include <nxos.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <stdlib.h>
#include <stdint.h>
#include <ctype.h>

#include <netsocket.h>

#define HOST_ADDR "192.168.0.200"

int TCP_ClientTest(void)
{
    NX_Error err;
    char buf[256];
    
    printf("CLIENT: hello!\n");

    int sock = NET_Socket(AF_INET, SOCK_STREAM, 0);
    printf("CLIENT: sock: %d\n", sock);
    if (sock < 0)
    {
        return -1;
    }
    struct sockaddr_in addr_in;

    addr_in.sin_family = AF_INET;
    addr_in.sin_port = htons(8888);
    addr_in.sin_addr.s_addr = inet_addr(HOST_ADDR);
    err = NX_Connect(sock, (struct sockaddr *)&addr_in, sizeof(addr_in));
    if (err != NX_EOK)
    {
        printf("CLIENT: exit on connect\n");

        NET_Close(sock);
        return 1;
    }

    int pkglen;

    strcpy(buf, "test data abc 1234!!!");
    pkglen = NET_Send(sock, buf, strlen(buf), 0);
    if (pkglen <= 0)
    {
        printf("CLIENT: send err: %d\n", pkglen);
        NET_Close(sock);
        return 1;
    }
    printf("CLIENT: send len %d success!\n", pkglen);
    
    memset(buf, 0, 256);
    /* do read/write */
    pkglen = NET_Recv(sock, buf, sizeof(buf), 0);
    if (pkglen <= 0)
    {
        printf("CLIENT: recv err: %d\n", pkglen);
        NET_Close(sock);
        return 1;
    }
    printf("CLIENT: recv len %d data: %s\n", pkglen, buf);
    NET_Shutdown(sock, SHUT_RD);
    NET_Close(sock);
    
    return 0;
}

#define USE_MULTI_THREAD 0

int HandleAccept(int conn)
{
  int pkglen;
  int err;
  char buf[256];
  
  printf("SERVER: connect %d\n", conn);

  unsigned long nrecv = -1;
  err = NET_Control(conn, NET_IONREAD, &nrecv);
  if (err != NX_EOK)
  {
      printf("SERVER: control failed!\n");
  }
  printf("SERVER: NET_IONREAD %d\n", nrecv);

  /* do read/write */
  pkglen = NET_Recv(conn, buf, sizeof(buf), 0);
  if (pkglen <= 0)
  {
      printf("SERVER: recv err: %d\n", pkglen);
      NET_Close(conn);
      return -1;
  }
  printf("SERVER: recv len %d data: %s\n", pkglen, buf);

  pkglen = NET_Send(conn, buf, pkglen, 0);
  if (pkglen <= 0)
  {
      printf("SERVER: send err: %d\n", pkglen);
      NET_Close(conn);
      return -1;
  }
  printf("SERVER: send len %d success!\n", pkglen);

  NET_Shutdown(conn, SHUT_WR);
  NET_Close(conn);
  return 0;
}

#if USE_MULTI_THREAD ==1
NX_U32 AcceptThread(void *arg)
{
  return HandleAccept((long)arg);
}

#endif

int TCP_ServerTest(void)
{
    NX_Error err;

    int sock = NET_Socket(AF_INET, SOCK_STREAM, 0);
    printf("SERVER: sock: %d\n", sock);
    if (sock < 0) {
        return -1;
    }
    struct sockaddr_in addr_in;
    struct sockaddr client;

    socklen_t client_len;

    addr_in.sin_family = AF_INET;
    addr_in.sin_port = htons(8888);
    addr_in.sin_addr.s_addr = inet_addr(HOST_ADDR);
    err = NET_Bind(sock, (struct sockaddr *)&addr_in, sizeof(struct sockaddr));
    if (err != NX_EOK)
    {
        printf("SERVER: exit on bind\n");
        NET_Close(sock);
        return 1;
    }
    err = NET_Listen(sock, 8);
    if (err != NX_EOK)
    {
        printf("SERVER: exit on listen\n");
        NET_Close(sock);
        return 1;
    }

    while (1)
    {
        printf("SERVER: accept ready: %d\n", sock);
        int conn = NET_Accept(sock, &client, &client_len);
        if (conn < 0)
        {
            printf("SERVER: accept err on sock: %d\n", conn);        
            // printf("SERVER: exit on accept\n");
            continue;
        }
        else
        {
            printf("SERVER: accept ok: %d\n", conn);
#if USE_MULTI_THREAD == 1
            NX_Solt child = NX_ThreadCreate(NX_NULL, HandleAccept, (void *)conn, 0);
            NX_Printf("server create child %d, connect %d\n", child, conn);
            // NX_ThreadWait(child, NX_NULL);
#else
            HandleAccept(conn);
#endif
        }
    }
    NET_Shutdown(sock, SHUT_RDWR);
    NET_Close(sock);
    return 0;
}

int UDP_ClientTest(void)
{
    char buf[256];
    
    printf("CLIENT: hello!\n");

    int sock = NET_Socket(AF_INET, SOCK_DGRAM, 0);
    printf("CLIENT: sock: %d\n", sock);
    if (sock < 0)
    {
        return -1;
    }
    struct sockaddr_in addr_in;
    struct sockaddr addr_server;
    socklen_t addr_len;

    addr_in.sin_family = AF_INET;
    addr_in.sin_port = htons(8888);
    addr_in.sin_addr.s_addr = inet_addr(HOST_ADDR);

    int pkglen;

    strcpy(buf, "test data abc 1234!!!");
    pkglen = NET_Sendto(sock, buf, strlen(buf), 0, (struct sockaddr *)&addr_in, sizeof(addr_in));
    if (pkglen <= 0)
    {
        printf("CLIENT: send err: %d\n", pkglen);
        NET_Close(sock);
        return 1;
    }
    printf("CLIENT: send len %d success!\n", pkglen);
    
    memset(buf, 0, 256);
    /* do read/write */
    pkglen = NET_Recvfrom(sock, buf, sizeof(buf), 0, &addr_server, &addr_len);
    if (pkglen <= 0)
    {
        printf("CLIENT: recv err: %d\n", pkglen);
        NET_Close(sock);
        return 1;
    }
    printf("CLIENT: recv len %d data: %s, addr len:%d\n", pkglen, buf, addr_len);
    NET_Close(sock);
    
    return 0;
}

int UDP_ServerTest(void)
{
    NX_Error err;
    char buf[256];

    int sock = NET_Socket(AF_INET, SOCK_DGRAM, 0);
    printf("SERVER: sock: %d\n", sock);
    if (sock < 0) {
        return -1;
    }
    struct sockaddr_in addr_in;
    struct sockaddr client;

    socklen_t client_len;

    addr_in.sin_family = AF_INET;
    addr_in.sin_port = htons(8888);
    addr_in.sin_addr.s_addr = inet_addr(HOST_ADDR);
    err = NET_Bind(sock, (struct sockaddr *)&addr_in, sizeof(struct sockaddr));
    if (err != NX_EOK)
    {
        printf("SERVER: exit on bind\n");
        NET_Close(sock);
        return 1;
    }

    while (1)
    {
        /* do read/write */
        int pkglen = NET_Recvfrom(sock, buf, sizeof(buf), 0, &client, &client_len);
        if (pkglen <= 0)
        {
            printf("SERVER: recv err: %d\n", pkglen);
            break;
        }
        printf("SERVER: recv len %d data: %s, addr len:%d\n", pkglen, buf, client_len);

        pkglen = NET_Sendto(sock, buf, pkglen, 0, &client, client_len);
        if (pkglen <= 0)
        {
            printf("SERVER: send err: %d\n", pkglen);
            break;
        }
        printf("SERVER: send len %d success!\n", pkglen);
    }
    NET_Shutdown(sock, SHUT_RDWR);
    NET_Close(sock);
    return 0;
}

static const char http_html_hdr[] = "HTTP/1.1 200 OK\r\nContent-type: text/html\r\n\r\n";
static const char http_index_html[] = "<html><head><title>Congrats!</title></head><body><h1>Welcome to our lwIP HTTP server!</h1><p>This is a small test page, served by httpserver-netconn.</body></html>";

int WebServerTest(void)
{
    NX_Error err;
    char buf[256];

    int sock = NET_Socket(AF_INET, SOCK_STREAM, 0);
    printf("SERVER: sock: %d\n", sock);
    if (sock < 0) {
        return -1;
    }
    struct sockaddr_in addr_in;
    struct sockaddr client;

    socklen_t client_len;

    addr_in.sin_family = AF_INET;
    addr_in.sin_port = htons(80);
    addr_in.sin_addr.s_addr = inet_addr(HOST_ADDR);
    err = NET_Bind(sock, (struct sockaddr *)&addr_in, sizeof(struct sockaddr));
    if (err != NX_EOK)
    {
        printf("SERVER: exit on bind\n");
        NET_Close(sock);
        return 1;
    }
    err = NET_Listen(sock, 8);
    if (err != NX_EOK)
    {
        printf("SERVER: exit on listen\n");
        NET_Close(sock);
        return 1;
    }

    while (1)
    {
        int conn = NET_Accept(sock, &client, &client_len);
        if (conn < 0)
        {
            printf("SERVER: accept err on sock: %d\n", sock);        
            printf("SERVER: exit on accept\n");
            break;
        }
        else
        {
            printf("SERVER: accept ok: %d\n", conn);
            int pkglen;

            /* do read/write */
            pkglen = NET_Recv(conn, buf, sizeof(buf), 0);
            if (pkglen <= 0)
            {
                printf("SERVER: recv err: %d\n", pkglen);
                NET_Close(conn);
                break;
            }
            printf("SERVER: recv len %d data: %s\n", pkglen, buf);

            /* check http */
            if (pkglen>=5 &&
                buf[0]=='G' &&
                buf[1]=='E' &&
                buf[2]=='T' &&
                buf[3]==' ' &&
                buf[4]=='/' ) {

                /* Send the HTML header
                        * subtract 1 from the size, since we don't send the \0 in the string
                        * NETCONN_NOCOPY: our data is const static, so no need to copy it
                */
                pkglen = NET_Send(conn, (void *)http_html_hdr, sizeof(http_html_hdr)-1, 0);
                if (pkglen <= 0)
                {
                    printf("SERVER: send err: %d\n", pkglen);
                    NET_Close(conn);
                    break;
                }
                printf("SERVER: send len %d success!\n", pkglen);

                /* Send our HTML page */
                pkglen = NET_Send(conn, (void *)http_index_html, sizeof(http_index_html)-1, 0);
                if (pkglen <= 0)
                {
                    printf("SERVER: send err: %d\n", pkglen);
                    NET_Close(conn);
                    break;
                }
                printf("SERVER: send len %d success!\n", pkglen);
            }

            NET_Shutdown(conn, SHUT_WR);
            NET_Close(conn);
        }
    }
    NET_Shutdown(sock, SHUT_RDWR);
    NET_Close(sock);
    return 0;
}

#define RECV_TIMEOUT_USEC 1000000

int SockoptTest(void)
{
    NX_Error err; 
    int sock = NET_Socket(AF_INET, SOCK_STREAM, 0);
    printf("SERVER: sock: %d\n", sock);
    if (sock < 0) {
        return -1;
    }
    struct timeval tv;
    tv.tv_sec = 0;
    tv.tv_usec = RECV_TIMEOUT_USEC;
    err = NET_Setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    printf("NET_Setsockopt: err %d\n", err);
    struct timeval tv2;
    socklen_t optlen = sizeof(tv2);
    err = NET_Getsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv2, &optlen);
    printf("NET_Getsockopt: err %d, tv2.tv_sec:%d, tv.tv_usec:%d\n", err, tv2.tv_sec, tv.tv_usec);

    NET_Close(sock);

    return 0;
}
#define MAGIC "1234567890"
#define MAGIC_LEN 11
#define MTU 1500

struct icmp_echo {
    // header
    uint8_t type;
    uint8_t code;
    uint16_t checksum;

    uint16_t ident;
    uint16_t seq;

    // data
    uint32_t sending_ts;
    char magic[MAGIC_LEN];
};

uint32_t get_timestamp()
{
    return NX_ClockGetMillisecond();
}

uint16_t calculate_checksum(unsigned char* buffer, int bytes)
{
    uint32_t checksum = 0;
    unsigned char* end = buffer + bytes;

    // odd bytes add last byte and reset end
    if (bytes % 2 == 1) {
        end = buffer + bytes - 1;
        checksum += (*end) << 8;
    }

    // add words of two bytes, one by one
    while (buffer < end) {
        checksum += buffer[0] << 8;
        checksum += buffer[1];
        buffer += 2;
    }

    // add carry if any
    uint32_t carray = checksum >> 16;
    while (carray) {
        checksum = (checksum & 0xffff) + carray;
        carray = checksum >> 16;
    }

    // negate it
    checksum = ~checksum;

    return checksum & 0xffff;
}

int send_echo_request(int sock, struct sockaddr_in* addr, int ident, int seq)
{
    // allocate memory for icmp packet
    struct icmp_echo icmp;
    memset(&icmp, 0, sizeof(icmp));

    // fill header files
    icmp.type = 8;
    icmp.code = 0;
    icmp.ident = htons(ident);
    icmp.seq = htons(seq);

    // fill magic string
    strncpy(icmp.magic, MAGIC, MAGIC_LEN);

    // fill sending timestamp
    icmp.sending_ts = get_timestamp();

    // calculate and fill checksum
    icmp.checksum = htons(
        calculate_checksum((unsigned char*)&icmp, sizeof(icmp))
    );

    // send it
    int bytes = NET_Sendto(sock, &icmp, sizeof(icmp), 0,
        (struct sockaddr*)addr, sizeof(*addr));
    if (bytes == -1) {
        return -1;
    }

    return 0;
}

int recv_echo_reply(int sock, int ident)
{
    // allocate buffer
    char buffer[MTU];
    struct sockaddr_in peer_addr;

    // receive another packet
    socklen_t addr_len = sizeof(peer_addr);
    int bytes = NET_Recvfrom(sock, buffer, sizeof(buffer), 0,
        (struct sockaddr*)&peer_addr, &addr_len);
    if (bytes == -1) {
        // normal return when timeout
        if (NX_ErrorGet() == NX_ETIMEOUT) {        
            return 0;
        }

        return -1;
    }

    // find icmp packet in ip packet
    struct icmp_echo* icmp = (struct icmp_echo*)(buffer + 20);

    // check type
    if (icmp->type != 0 || icmp->code != 0) {
        return 0;
    }

    // match identifier
    if (ntohs(icmp->ident) != ident) {
        return 0;
    }

    // print info
    printf("%s seq=%d %dms\n",
        inet_ntoa(peer_addr.sin_addr),
        ntohs(icmp->seq),
        (get_timestamp() - icmp->sending_ts)
    );

    return 0;
}

int PingTest(const char *ip)
{
    // for store destination address
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));

    // fill address, set port to 0
    addr.sin_family = AF_INET;
    addr.sin_port = 0;
    if (inet_aton(ip, (struct in_addr*)&addr.sin_addr.s_addr) == 0) {
        return -1;
    };

    printf("ping: %s\n", ip);

    // create raw socket for icmp protocol
    int sock = NET_Socket(AF_INET, SOCK_RAW, IP_PROTO_ICMP);
    if (sock == -1) {
        return -1;
    }

    // set socket timeout option
    struct timeval tv;
    tv.tv_sec = 0;
    tv.tv_usec = RECV_TIMEOUT_USEC;
    int ret;
    NX_Error err = NET_Setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    if (err != NX_EOK) {
        return -1;
    }

    uint32_t next_ts = get_timestamp();
    int ident = NX_ThreadGetCurrentId();
    int seq = 1;

    for (;;) {
        // time to send another packet
        if (get_timestamp() >= next_ts) {
            // send it
            ret = send_echo_request(sock, &addr, ident, seq);
            if (ret == -1) {
                perror("Send failed");
            }

            // update next sendint timestamp to one second later
            next_ts += 1000;
            // increase sequence number
            seq += 1;
        }

        // try to receive and print reply
        ret = recv_echo_reply(sock, ident);
        if (ret == -1) {
            perror("Receive failed");
        }
    }

    return 0;
}

int main(int argc, char **argv)
{
#if 1
    if (argc > 1)
    {
        printf("hello, socket test client!\n");
#if 1
        return TCP_ClientTest();
#else
        return UDP_ClientTest();
#endif
    }
    else
    {
        printf("hello, socket test server!\n");
#if 1
        return TCP_ServerTest();
#else
        return UDP_ServerTest();
#endif
    }
#else
    // return WebServerTest();
    // return SockoptTest();
    if (argc < 2) {
        printf("usage: \n\tping x.x.x.x\n");
    } else {
        return PingTest(argv[1]);
    }
#endif
}
